package com.hexinfo.dmpro.data.transfer.common.common.manager;

import cn.hutool.core.date.DateUtil;
import com.hexinfo.dmpro.data.transfer.common.common.session.AgentSession;
import com.hexinfo.dmpro.data.transfer.common.model.AgentInfo;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class AgentManager {

    Map<String, AgentSession> agentSessionMap = new ConcurrentHashMap<>();

    /*
     * @Method: register <br>
     * @Param: [agentInfo] <br>
     * @Return: void <br>
     * @Description：AGENT注册<br>
     * @Author： wz.li<br>
     * @Date： 2023/11/13 14:31 <br>
     * @Version： V2.0.2<br>
     */
    public void register(AgentInfo agentInfo) {
        unregister(agentInfo);
        agentSessionMap.put(agentInfo.getAgentKey(), new AgentSession(agentInfo));
    }

    /**
     * AGENT 更新心跳信息
     *
     * @param agentInfo
     */
    public void updateAgentSession(AgentInfo agentInfo) {
        AgentSession agentSession = agentSessionMap.get(agentInfo.getAgentKey());
        if (agentSession == null) {
            register(agentInfo);
        } else {
            agentSession.setMaxPoolSize(agentInfo.getMaxPoolSize());
            agentSession.setLastHeartbeatTime(DateUtil.now());
        }
    }

    /*
     * @Method: unregister <br>
     * @Param: [agentInfo] <br>
     * @Return: void <br>
     * @Description：Agent注销<br>
     * @Author： wz.li<br>
     * @Date： 2023/11/13 14:32 <br>
     * @Version： V2.0.2<br>
     */
    public void unregister(AgentInfo agentInfo) {
        agentSessionMap.remove(agentInfo.getAgentKey());
    }

    public String getAvailableAgent() {
        List<AgentSession> agentSessionList = new ArrayList<>(agentSessionMap.values());
        Collections.sort(agentSessionList, Comparator.comparing(AgentSession::getAgentIdleNum).reversed());
        for (AgentSession agentSession : agentSessionList) {
            if (agentSession.isAvailable()) {
                return agentSession.getAgentKey();
            }
        }
        return null;
    }

    public boolean addExecuteTask(String agentKey) {
        AgentSession agentSession = agentSessionMap.get(agentKey);
        return agentSession.addExecuteTask();
    }

    public boolean removeExecuteTask(String agentKey) {
        AgentSession agentSession = agentSessionMap.get(agentKey);
        return agentSession.delExecuteTask();
    }

    public List<AgentSession> agentSessionList() {
        return new ArrayList<>(agentSessionMap.values());
    }

}
